#include <iostream>
#include <stdlib.h>
#include <random>
#include <cmath>

#include "gtest/gtest.h"

#include "kpextVkContext.h"
#include "KpExtGP.h"
#include "kpextOpImgCMapF32C1toU8C4.h"

#include "kpextCtxGlobal.h"

using namespace std;

#define INPUT_MAX_V 65535
#define COLOR_MAP_LEN 1024
typedef struct{
    float r;
    float g;
    float b;
    float a;
}FRGBA_SZ;

std::array<uint8_t, 4> colormap(std::shared_ptr<float> map, float value){
    std::array<uint8_t, 4> ret;
    uint32_t index;
    value  = value / INPUT_MAX_V;
    value  = value * (COLOR_MAP_LEN - 1.0f) + 0.5f;
    value  = value * float(value >= 0);
    index = static_cast<uint32_t>(value);
    index = index * (index < 1024U) + 1023U * (index > 1023U);

    //  映射
    auto ptr = reinterpret_cast<FRGBA_SZ *>(map.get());
    auto color = ptr[index];

    //  RGBA转换到0~255
    index = static_cast<uint32_t>(color.r + 0.5f);
    index = index * (index < 256U) + 255U * (index > 255U);
    ret[0] = static_cast<uint8_t>(index);

    index = static_cast<uint32_t>(color.g + 0.5f);
    index = index * (index < 256U) + 255U * (index > 255U);
    ret[1] = static_cast<uint8_t>(index);

    index = static_cast<uint32_t>(color.b + 0.5f);
    index = index * (index < 256U) + 255U * (index > 255U);
    ret[2] = static_cast<uint8_t>(index);

    index = static_cast<uint32_t>(color.a + 0.5f);
    index = index * (index < 256U) + 255U * (index > 255U);
    ret[3] = static_cast<uint8_t>(index);

    return ret;
}

TEST(kpextFuncTest, OpImgCMapF32C1U8C4Test){
    ASSERT_NE(nullptr, GvkContextPtr);

    int cols = 1024, rows = 1000;

    auto cmdbuf = GvkContextPtr->cmdbuf();

    size_t size = static_cast<size_t>(cols * rows);
    auto stagebuf1 = GvkContextPtr->stagingBuf(size * 4);
    auto stagebuf2 = GvkContextPtr->stagingBuf(size * 4, false, true);
    auto inbuf = GvkContextPtr->deviceBuf(size * 4);
    auto outbuf = GvkContextPtr->deviceBuf(size * 4);
    auto op = std::make_shared<kpext::OpImgCMapF32C1U8C4>(
                GvkContextPtr
                );

    ASSERT_NE(nullptr, cmdbuf);
    ASSERT_NE(nullptr, stagebuf1);
    ASSERT_NE(nullptr, stagebuf2);
    ASSERT_NE(nullptr, inbuf);
    ASSERT_NE(nullptr, outbuf);
    ASSERT_NE(nullptr, op);


    {
        std::random_device r;

        // Choose a random mean between 0 and 65535
        std::default_random_engine e1(r());
        std::uniform_int_distribution<int32_t> uniform_dist(-10000, 70000);

        auto data = stagebuf1->data<float>();
        for (size_t i = 0; i < size; ++i) {
            data[i] = static_cast<float>(uniform_dist(e1));//static_cast<float>(i);
        }
        stagebuf1->unmap();
    }


    kpext::CmdBaseBufferCpy bufcpy;
    kpext::CmdBaseBufferSet bufset;
    kpext::CmdBuffersMemBarrier barrierTx;
    kpext::CmdBuffersMemBarrier barrierRx;

    barrierTx.prepare(
                {inbuf, outbuf},
                vk::PipelineStageFlagBits::eTransfer,
                vk::PipelineStageFlagBits::eComputeShader,
                vk::AccessFlagBits::eTransferWrite,
                vk::AccessFlagBits::eShaderRead
                );
    barrierRx.prepare(
                {inbuf, outbuf},
                vk::PipelineStageFlagBits::eComputeShader,
                vk::PipelineStageFlagBits::eTransfer,
                vk::AccessFlagBits::eShaderWrite,
                vk::AccessFlagBits::eTransferRead
                );

    // 纯灰度颜色映射测试
    op->rebuild(cols, rows,
                0,
                inbuf, outbuf,
                cmdbuf
                );

    cmdbuf->begin();
    bufcpy.prepare(inbuf, stagebuf1, size * 4);
    cmdbuf->record(&bufcpy);
    cmdbuf->record(&barrierTx);

    cmdbuf->record(op.get());

    bufcpy.prepare(stagebuf2, outbuf);
    cmdbuf->record(&barrierRx);
    cmdbuf->record(&bufcpy);

    cmdbuf->submit();
    cmdbuf->waitExec();

    {
        auto cmap = kpext::OpImgCMapF32C1U8C4::getColorMapClone(0);
        auto p1 = stagebuf1->data<float>();
        auto p2 = stagebuf2->data<uint8_t>();
        stagebuf2->invalidateMapped();
//            SPDLOG_INFO(__FUNCTION__ " input: [{}] ", stagebuf1->vector<float>());
//            SPDLOG_INFO(__FUNCTION__ " output: [{}] ", stagebuf2->vector<uint8_t>());

        int maxerr = 0;
        for (int r = 0; r < rows; ++r) {
            for (int c = 0; c < cols; ++c) {
                int idx = c + r * cols;
                float input =  p1[idx];
                std::array<uint8_t, 4> output;
                if(cmap){
                    output = colormap(cmap, input);
                }else{
                    for (auto & o : output) {
                        o = static_cast<uint8_t>(input);
                    }
                }

                idx = idx * 4;
                for (int i = 0; i < 4; ++i) {
                    int d = p2[idx + i];
                    d = d - output[static_cast<size_t>(i)];
                    if(d < 0) d = -d;
                    if(d > maxerr){
                        maxerr = d;
                    }
                }
            }
//            auto sptr = p2 + r * cols * 4;
//            std::vector<uint8_t> rowdisp(sptr, sptr + cols * 4);
//            SPDLOG_INFO(" {} ", rowdisp);
        }
        stagebuf1->unmap();
        stagebuf2->unmap();
        std::cout << "max error = " << maxerr << std::endl;
        EXPECT_LE(maxerr, 2);
    }

    // 自定义颜色映射测试
    {
        std::vector<float> cmapdata{
            0,        0,  128,    0,  128,
            1024,   128,    0,  128,    0,
            65535,  255,  128,  255,  128,
        };
        kpext::OpImgCMapF32C1U8C4::setColorMap(1, cmapdata);
//        auto cmap = kpext::OpImgCMapF32C1U8C4::getColorMapClone(1);
//        SPDLOG_INFO("ColorMap[{}] : ", 1);
//        for (int r = 0; r < 256; ++r) {
//            float * ptr = cmap.get() + r * 4 * 4;
//            std::vector<float> pixel(ptr, ptr + 4 * 4);
//            std::vector<uint8_t> disp;
//            disp.resize(pixel.size());
//            std::transform(pixel.begin(), pixel.end(),
//                           disp.begin(),
//                           [](float v){
//                return static_cast<uint8_t>(v + 0.5f);
//            }
//            );
//            SPDLOG_INFO("{}", disp);
//        }
    }

    op->rebuild(cols, rows,
                1,
                inbuf, outbuf,
                cmdbuf
                );

    cmdbuf->begin();

    cmdbuf->record(op.get());

    bufcpy.prepare(stagebuf2, outbuf);
    cmdbuf->record(&barrierRx);
    cmdbuf->record(&bufcpy);

    cmdbuf->submit();
    cmdbuf->waitExec();

    {
        auto cmap = kpext::OpImgCMapF32C1U8C4::getColorMapClone(1);
        auto p1 = stagebuf1->data<float>();
        auto p2 = stagebuf2->data<uint8_t>();
        stagebuf2->invalidateMapped();
//            SPDLOG_INFO(__FUNCTION__ " input: [{}] ", stagebuf1->vector<float>());
//            SPDLOG_INFO(__FUNCTION__ " output: [{}] ", stagebuf2->vector<uint8_t>());

        int maxerr = 0;
        for (int r = 0; r < rows; ++r) {
            for (int c = 0; c < cols; ++c) {
                int idx = c + r * cols;
                float input =  p1[idx];
                std::array<uint8_t, 4> output;
                if(cmap){
                    output = colormap(cmap, input);
                }else{
                    for (auto & o : output) {
                        o = static_cast<uint8_t>(input);
                    }
                }

                idx = idx * 4;
                for (int i = 0; i < 4; ++i) {
                    int d = p2[idx + i];
                    d = d - output[static_cast<size_t>(i)];
                    if(d < 0) d = -d;
                    if(d > maxerr){
                        maxerr = d;
                    }
                }
            }
//            auto sptr = p2 + r * cols * 4;
//            std::vector<uint8_t> rowdisp(sptr, sptr + cols * 4);
//            SPDLOG_INFO(" {} ", rowdisp);
        }
        stagebuf1->unmap();
        stagebuf2->unmap();
        std::cout << "max error = " << maxerr << std::endl;
        EXPECT_LE(maxerr, 2);
    }

}


GTEST_API_ int main(int argc, char **argv) {
    printf("Running main() from %s\n", __FILE__);
    //    kpextSetLoggerLevel(KpExtLOG_LEVEL_DEBUG);
    kpextSetLoggerLevel(KpExtLOG_LEVEL_INFO);
    //  kpextSetLoggerPattern("[%D %H:%M:%S.%e] [%^%l%$] %v");
    kpextSetLoggerPattern("   [SPDLOG_%l] %v");

    testing::InitGoogleTest(&argc, argv);

    testing::Environment* const foo_env =
            testing::AddGlobalTestEnvironment(new KpExtCtxEnv);
    UNUSE_VAR(foo_env);

    return RUN_ALL_TESTS();
}
